const fs = require('fs');
const path = require('path');

// 从命令行参数获取 CUDA 版本和 packageManager
const cudaVersion = process.argv[2];
const packageManager = process.argv[3] || 'pip'; // 默认为 pip

// 配置文件路径
const CONFIG_FILE = `pytorch_${packageManager}_versions.json`;

// 获取最接近的 CUDA 版本对应的 PyTorch 安装命令
function loadPytorchVersions() {
    const filePath = path.join(__dirname, CONFIG_FILE);
    const rawData = fs.readFileSync(filePath);
    return JSON.parse(rawData);
}

function getClosestVersionCommand(cudaVersion) {
    const pytorchVersions = loadPytorchVersions();
    const availableVersions = Object.keys(pytorchVersions).filter((key) => key.startsWith('CUDA'));
    const targetMajor = parseInt(cudaVersion.split('.')[0], 10);
    const targetMinor = parseFloat(cudaVersion);

    const matchingVersions = availableVersions.filter((v) => parseInt(v.split(' ')[1].split('.')[0], 10) === targetMajor);

    if (!matchingVersions.length) {
        return pytorchVersions['CPU Only'][0];
    }
    const closest = matchingVersions.reduce((prev, curr) => {
        const currMinor = parseFloat(curr.split(' ')[1]);
        return Math.abs(currMinor - targetMinor) < Math.abs(parseFloat(prev.split(' ')[1]) - targetMinor) ? curr : prev;
    });
    const command = pytorchVersions[closest][0] || pytorchVersions['CPU Only'][0];
    return command;
}

if (!cudaVersion) {
    console.error('Usage: node get_pytorch_command.js <cuda_version> [package_manager]');
    process.exit(1);
}

// 获取并输出安装命令
const installCommand = getClosestVersionCommand(cudaVersion, packageManager);
console.log(installCommand);